import numpy as np
import tensorflow as tf
from python_ai.common.xcommon import *


# iterator
def data_iter(x, y, batch_size=8):
    xlen = len(x)
    a = np.random.permutation(xlen).tolist()
    a.sort()  # test
    total_iters = int(np.ceil(xlen / batch_size))
    print(total_iters)
    for i in range(total_iters):
        print(i)
        idx = a[i*batch_size:(i+1)*batch_size]
        print(idx)
        yield tf.gather(x, idx), tf.gather(y, idx)


n = 20
x = tf.range(n)
y = 2 * tf.range(n)
batch_size = 8
total_iters = int(np.ceil(n / batch_size))
generator = data_iter(x, y, batch_size=batch_size)
for i in range(total_iters):
    sep(i)
    bx, by = next(generator)
